#include "VulkanApplication.h"
#include "BufferManager.h"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

const static float PI = 3.141592653f;
struct Matrix
{
    float m[16];

    void makeIdentity()
    {
        for (int i = 0; i < 16; i++) m[i] = 0.0f;
        m[0 * 4 + 0] = 1.0f;
        m[1 * 4 + 1] = 1.0f;
        m[2 * 4 + 2] = 1.0f;
        m[3 * 4 + 3] = 1.0f;
    }

    void makeRotate(float angle, float x, float y, float z)
    {
        float mag = sqrt(x * x + y * y + z * z);
        float sinAngle = sin(angle * PI / 180.0f);
        float cosAngle = cos(angle * PI / 180.0f);
        float xx, yy, zz, xy, yz, zx, xs, ys, zs, oneMinusCos;
        if (mag <= 0.0f) return;
        makeIdentity();

        x /= mag; y /= mag; z /= mag;
        xx = x * x; yy = y * y; zz = z * z;
        xy = x * y; yz = y * z; zx = z * x;
        xs = x * sinAngle; ys = y * sinAngle; zs = z * sinAngle;
        oneMinusCos = 1.0f - cosAngle;

        m[0 * 4 + 0] = (oneMinusCos * xx) + cosAngle;
        m[0 * 4 + 1] = (oneMinusCos * xy) - zs;
        m[0 * 4 + 2] = (oneMinusCos * zx) + ys;
        m[0 * 4 + 3] = 0.0f;
        m[1 * 4 + 0] = (oneMinusCos * xy) + zs;
        m[1 * 4 + 1] = (oneMinusCos * yy) + cosAngle;
        m[1 * 4 + 2] = (oneMinusCos * yz) - xs;
        m[1 * 4 + 3] = 0.0f;
        m[2 * 4 + 0] = (oneMinusCos * zx) - ys;
        m[2 * 4 + 1] = (oneMinusCos * yz) + xs;
        m[2 * 4 + 2] = (oneMinusCos * zz) + cosAngle;
        m[2 * 4 + 3] = 0.0f;
    }

    void makeOrtho(float l, float r, float b, float t, float zn, float zf)
    {
        makeIdentity();
        m[0 * 4 + 0] = 2.0f / (r - l);
        m[1 * 4 + 1] = 2.0f / (b - t);
        m[2 * 4 + 2] = 1.0f / (zn - zf);
        m[3 * 4 + 0] = -(r + l) / (r - l);
        m[3 * 4 + 1] = -(b + t) / (b - t);
        m[3 * 4 + 2] = zn / (zn - zf);
    }
};

int main(int argc, char** argv)
{
    VulkanApplication app("c12_uniforms", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create uniform desciption layout
    VkDescriptorSetLayoutBinding uboLayoutBinding
    {
        0,  // binding
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        1,  // descriptorCount
        VK_SHADER_STAGE_VERTEX_BIT,
        nullptr
    };

    VkDescriptorSetLayoutCreateInfo layoutInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        nullptr, 0,
        1, &uboLayoutBinding
    };

    VkDescriptorSetLayout descSetLayout = nullptr;
    VkResult result = vkCreateDescriptorSetLayout(app.getDevice(), &layoutInfo, nullptr, &descSetLayout);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description-set layout.\n");
        return 1;
    }

    // Create uniform desciption pool
    VkDescriptorPoolSize poolSize
    { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 };

    VkDescriptorPoolCreateInfo poolInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        nullptr, 0,
        1, 1, &poolSize
    };

    VkDescriptorPool descriptorPool = nullptr;
    result = vkCreateDescriptorPool(app.getDevice(), &poolInfo, nullptr, &descriptorPool);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description pool.\n");
        return 1;
    }

    // Create uniform description sets
    VkDescriptorSetAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        nullptr,
        descriptorPool,
        1, &descSetLayout
    };

    VkDescriptorSet descriptorSet = nullptr;
    result = vkAllocateDescriptorSets(app.getDevice(), &allocInfo, &descriptorSet);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description set.\n");
        return 1;
    }

    // Create pipeline, vertices and indices
    BufferManager manager(&app);
    if (!manager.prepare("draw_and_rotate.vert", "draw_and_rotate.frag", descSetLayout))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Create vertex & index buffer
    std::vector<Vertex> vertices;
    vertices.push_back({ -0.5f, -0.5f, 0.0f, 0.0f, 1.0f, 0.0f });
    vertices.push_back({ 0.5f, -0.5f, 0.0f, 0.0f, 0.0f, 1.0f });
    vertices.push_back({ 0.5f, 0.5f, 0.0f, 1.0f, 0.0f, 0.0f });
    vertices.push_back({ -0.5f, 0.5f, 0.0f, 1.0f, 1.0f, 0.0f });

    std::vector<unsigned int> indices;
    indices.push_back(0); indices.push_back(1); indices.push_back(2);
    indices.push_back(0); indices.push_back(2); indices.push_back(3);

    VkBuffer vertexBuffer = manager.createVertexBuffer("vertices", vertices);
    VkBuffer indexBuffer = manager.createIndexBuffer("indices", indices);

    // Create UBO buffer
    struct UniformBufferObject
    {
        Matrix modelview;
        Matrix projection;
    };

    float aspectRatio = (float)1024 / (float)768;
    UniformBufferObject ubo;
    ubo.modelview.makeRotate(30.0f, 0.0f, 0.0f, 1.0f);
    ubo.projection.makeOrtho(-aspectRatio, aspectRatio, -1.0f, 1.0f, -10.0f, 10.0f);

    size_t uboDataSize = sizeof(UniformBufferObject);
    VkBuffer uniformBuffer = nullptr;
    VkDeviceMemory uniformBufferMemory = nullptr;
    if (!manager.createBuffer(uniformBuffer, uboDataSize, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT))
    {
        printf("Failed to create uniform buffer.\n");
        return 1;
    }

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(app.getDevice(), uniformBuffer, &memRequirements);
    if (!manager.createBufferMemory(uniformBufferMemory, memRequirements,
                                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
    {
        printf("Failed to create uniform buffer memory.\n");
        return 1;
    }

    void* uboData = nullptr;
    vkBindBufferMemory(app.getDevice(), uniformBuffer, uniformBufferMemory, 0);
    vkMapMemory(app.getDevice(), uniformBufferMemory, 0, uboDataSize, 0, &uboData);
    memcpy(uboData, &ubo, uboDataSize);
    vkUnmapMemory(app.getDevice(), uniformBufferMemory);

    // Update description set
    VkDescriptorBufferInfo bufferInfo
    { uniformBuffer, 0, uboDataSize };

    VkWriteDescriptorSet descriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        descriptorSet, 0,  // dstSet, dstBinding
        0, 1,              // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        nullptr, &bufferInfo, nullptr
    };
    vkUpdateDescriptorSets(app.getDevice(), 1, &descriptorWrite, 0, nullptr);
    
    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
                                                waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        app.beginRendering(imageIndex, width, height);
        {
            VkBuffer vertexBuffers[] = { vertexBuffer };
            VkDeviceSize offsets[] = { 0 };

            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                    app.getPipelineLayout(), 0, 1, &descriptorSet, 0, nullptr);

            vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
            vkCmdDrawIndexed(commandBuffer, (uint32_t)indices.size(), 1, 0, 0, 0);
        }
        app.endRendering();

        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkDeviceWaitIdle(app.getDevice());
    vkDestroyDescriptorPool(app.getDevice(), descriptorPool, nullptr);
    vkDestroyDescriptorSetLayout(app.getDevice(), descSetLayout, nullptr);
    vkDestroyBuffer(app.getDevice(), uniformBuffer, nullptr);
    vkFreeMemory(app.getDevice(), uniformBufferMemory, nullptr);
    manager.cleanup();
    return 0;
}